Spark_url = "ws://spark-api.xf-yun.com/v3.1/chat"  # v3.0环境的地址
domain = "generalv3"    # v3.0版本
import _thread as thread
import base64,json,datetime,hashlib,hmac,websocket,ssl,time
from urllib.parse import urlparse
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
from wsgiref.handlers import format_date_time
answer = ""
def print_centered(text, width=80):#根据给定的宽度，居中打印文本.
    spaces = (width - len(text)) // 2
    centered_text = ' ' * spaces + text + ' ' * spaces
    print(centered_text)

class Ws_Param(object):
    def __init__(self, APPID, APIKey, APISecret, Spark_url):# 初始化
        self.APPID = APPID
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.host = urlparse(Spark_url).netloc
        self.path = urlparse(Spark_url).path
        self.Spark_url = Spark_url
    # 生成url
    def create_url(self):
        # 生成RFC1123格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))
        # 拼接字符串
        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"
        # 进行hmac-sha256进行加密
        signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
                                 digestmod=hashlib.sha256).digest()
        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')
        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')
        # 将请求的鉴权参数组合为字典
        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        # 拼接鉴权参数，生成url
        url = self.Spark_url + '?' + urlencode(v)
        # 此处打印出建立连接时候的url,参考本demo的时候可取消上方打印的注释，比对相同参数时生成的url与自己代码生成的url是否一致
        return url
# 收到websocket错误的处理
def on_error(ws, error):
    print("### error:", error)
# 收到websocket关闭的处理
def on_close(ws,one,two):
    print(" ")
# 收到websocket连接建立的处理
def on_open(ws):
    thread.start_new_thread(run, (ws,))
def run(ws, *args):
    data = json.dumps(gen_params(appid=ws.appid, domain= ws.domain,question=ws.question))
    ws.send(data)
# 收到websocket消息的处理
def on_message(ws, message):
    data = json.loads(message)
    code = data['header']['code']
    if code != 0:
        print(f'请求错误: {code}, {data}')
        ws.close()
    else:
        choices = data["payload"]["choices"]
        status = choices["status"]
        content = choices["text"][0]["content"]
        print(content,end ="")
        global answer
        answer += content
        if status == 2:
            ws.close()
def gen_params(appid, domain,question):
    #通过appid和用户的提问来生成请参数
    data = {
        "header": {
            "app_id": appid,
            "uid": "1234"
        },
        "parameter": {
            "chat": {
                "domain": domain,
                "temperature": 0.5,
                "max_tokens": 2048
            }
        },
        "payload": {
            "message": {
                "text": question
            }
        }
    }
    return data
def SparkApimain(appid, api_key, api_secret, Spark_url,domain, question):
    wsParam = Ws_Param(appid, api_key, api_secret, Spark_url)
    websocket.enableTrace(False)
    wsUrl = wsParam.create_url()
    ws = websocket.WebSocketApp(wsUrl, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open)
    ws.appid = appid
    ws.question = question
    ws.domain = domain
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
def read_config(filename="config.json"):
    try:
        with open(filename, "r") as file:
            config_data = json.load(file)
        return config_data
    except FileNotFoundError:
        print(f"错误: {filename} 无法找到.")
        return None
if __name__ == '__main__':
    print_centered("欢迎使用讯飞星火认知大模型3.0软件版", width=80)
    print_centered("本程序由裘炯涛老师（hznupeter）开发", width=80)
    print_centered("本项目开源地址：https://gitee.com/hznupeter/spark-desk.git", width=80)
    print_centered("=========接下来你就可以跟星火认知大模型对话啦=========", width=80)
    config = read_config()
    if config:
        appid = config.get("appid", "")
        api_secret = config.get("api_secret", "")
        api_key = config.get("api_key", "")
        jiange=config.get("jiange", "")
        if not all([appid, api_secret, api_key]):
            print("错误：不正确的配置文件，请提供正确的appid, api_secret和 api_key.")
            exit()
        text = []
        def get_text(role, content):
            jsoncon = {"role": role, "content": content}
            text.append(jsoncon)
            return text
        def get_length(text):
            length = 0
            for content in text:
                temp = content["content"]
                leng = len(temp)
                length += leng
            return length
        def check_len(text):
            if get_length(text)<5:
                return "您输入的内容太少，请重新输入"
            while get_length(text) > 8000:
                del text[0]
            return text
        text.clear()  # Corrected the function call
        time_since_last_question = 0
        while True:  # Infinite loop
            user_input = input("\n我:")
            current_time = time.time()
            if current_time - time_since_last_question >= jiange:#两次提问间隔要的大于jiange
                question = check_len(get_text("user", user_input))
                if isinstance(question, str):
                    print(question)
                    continue
                spark_api_answer = ""
                print("星火:", end="")
                SparkApimain(appid, api_key, api_secret, Spark_url, domain, question)
                text = get_text("assistant", spark_api_answer)  # Store the assistant's response
                time_since_last_question = time.time()
            else:#如果提问间隔太小，则提示
                remaining_time = int(jiange - (current_time - time_since_last_question))
                print(f"请等待 {remaining_time} 秒后再提问。")
                time.sleep(1)
    else:
        print("config.json 文件丢失，程序意外退出！")